package com.hdu.hystrix;

import com.hdu.context.FullLinkContextHolder;
import com.netflix.hystrix.HystrixThreadPoolKey;
import com.netflix.hystrix.HystrixThreadPoolProperties;
import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariable;
import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariableLifecycle;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ThreadPoolExecutor;

public class ThreadLocalAwareStrategy extends HystrixConcurrencyStrategy {


    private final HystrixConcurrencyStrategy existingStrategy;

    public ThreadLocalAwareStrategy(HystrixConcurrencyStrategy existingStrategy) {
        this.existingStrategy = existingStrategy;
    }

    @Override
    public BlockingQueue<Runnable> getBlockingQueue(int maxQueueSize) {
        return existingStrategy != null
                ? existingStrategy.getBlockingQueue(maxQueueSize)
                : super.getBlockingQueue(maxQueueSize);
    }

    @Override
    public <T> HystrixRequestVariable<T> getRequestVariable(HystrixRequestVariableLifecycle<T> rv) {
        return existingStrategy != null
                ? existingStrategy.getRequestVariable(rv)
                : super.getRequestVariable(rv);
    }

    @Override
    public ThreadPoolExecutor getThreadPool(HystrixThreadPoolKey threadPoolKey, HystrixThreadPoolProperties threadPoolProperties) {
        return existingStrategy != null
                ? existingStrategy.getThreadPool(threadPoolKey, threadPoolProperties)
                : super.getThreadPool(threadPoolKey, threadPoolProperties);
    }


    /**
     * 传递全链路压测上下文
     * @param callable
     * @return
     * @param <T>
     */
    @Override
    public <T> Callable<T> wrapCallable(Callable<T> callable) {
        return existingStrategy != null
                ? existingStrategy
                .wrapCallable(new DelegatingFullLinkContextCallable<>(callable, FullLinkContextHolder.getFullLinkContext()))
                : super.wrapCallable(new DelegatingFullLinkContextCallable<>(callable, FullLinkContextHolder.getFullLinkContext()));
    }
}
